{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"align-center\">\n",
    "<a href=\"https://oumi.ai/\"><img src=\"https://oumi.ai/docs/en/latest/_static/logo/header_logo.png\" height=\"200\"></a>\n",
    "\n",
    "[![Documentation](https://img.shields.io/badge/Documentation-latest-blue.svg)](https://oumi.ai/docs/en/latest/index.html)\n",
    "[![Discord](https://img.shields.io/discord/1286348126797430814?label=Discord)](https://discord.gg/oumi)\n",
    "[![GitHub Repo stars](https://img.shields.io/github/stars/oumi-ai/oumi)](https://github.com/oumi-ai/oumi)\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/oumi-ai/oumi/blob/main/notebooks/Oumi - Train a Letter Counting Model using GRPO.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "</div>\n",
    "\n",
    "👋 Welcome to Open Universal Machine Intelligence (Oumi)!\n",
    "\n",
    "🚀 Oumi is a fully open-source platform that streamlines the entire lifecycle of foundation models - from [data preparation](https://oumi.ai/docs/en/latest/resources/datasets/datasets.html) and [training](https://oumi.ai/docs/en/latest/user_guides/train/train.html) to [evaluation](https://oumi.ai/docs/en/latest/user_guides/evaluate/evaluate.html) and [deployment](https://oumi.ai/docs/en/latest/user_guides/launch/launch.html). Whether you're developing on a laptop, launching large scale experiments on a cluster, or deploying models in production, Oumi provides the tools and workflows you need.\n",
    "\n",
    "🤝 Make sure to join our [Discord community](https://discord.gg/oumi) to get help, share your experiences, and contribute to the project! If you are interested in joining one of the community's open-science efforts, check out our [open collaboration](https://oumi.ai/community) page.\n",
    "\n",
    "⭐ If you like Oumi and you would like to support it, please give it a star on [GitHub](https://github.com/oumi-ai/oumi)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a Letter Counting Model using GRPO\n",
    "\n",
    "Welcome to Oumi! In this tutorial notebook, we're going to fine-tune an LLM using Group Relative Policy Optimization (GRPO), a reinforcement learning algorithm. But first, a little (recent) history lesson --\n",
    "\n",
    "In June 2024, a user discovered that ChatGPT had a little problem -- it couldn't correctly answer a simple question, [\"How Many R’s Are in the Word Strawberry?\"](https://community.openai.com/t/incorrect-count-of-r-characters-in-the-word-strawberry/829618/2)\n",
    "\n",
    "Because of the way LLMs tokenize input strings, counting letters can be pretty tough for them! Fortunately, you (and Oumi!) are here to help.\n",
    "\n",
    "Below, we show you how to employ a custom evaluation function to evaluate popular models on the task of counting letters in words. Then, we will align Llama 3.2 3B to improve its performance on this task.\n",
    "\n",
    "This notebook includes cell outputs, but some irrelevant outputs (ex. install lines, warnings) are modified/removed for readability."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "### Machine Requirements\n",
    "\n",
    "This notebook runs both model evaluation and GRPO training, which require 8GB and 40GB VRAM, respectively.\n",
    "\n",
    "❗**NOTICE:** If you're running this notebook on Colab using a T4 GPU, it's not possible to run training due to memory requirements. To run evaluation, some adjustments need to be made as vLLM doesn't support T4 GPUs. This will be explained in the evaluation section.\n",
    "\n",
    "If your local machine cannot run this notebook, you can instead run this notebook on a cloud platform. The following demonstrates how to open a VSCode instance backed by a GCP node with 4 A100 GPUs, from which the notebook can be run. It is possible to run this notebook on just 1 GPU, but you will need make some adjustments to training parameters, which will be explained in the training section.\n",
    "\n",
    "```bash\n",
    "# Run on your local machine\n",
    "gcloud auth application-default login  # Authenticate with GCP\n",
    "make gcpcode ARGS=\"--resources.accelerators A100:4\"\n",
    "```\n",
    "\n",
    "### Oumi Installation\n",
    "\n",
    "First, let's install Oumi and vLLM (part of the `gpu` optional dependencies). You can find more detailed instructions about Oumi installation [here](https://oumi.ai/docs/en/latest/get_started/installation.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install oumi[gpu]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remote API Access\n",
    "\n",
    "As part of this notebook, you can evaluate frontier models from Open AI, Google, Anthropic, and Meta on the letter counting task. If you want to evaluate any of these models, set the corresponding fields below. The code is commented out by default to avoid any accidental overwriting of existing variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# os.environ[\"OPENAI_API_KEY\"] = \"\"  # Set your OpenAI API key here.\n",
    "# os.environ[\"GEMINI_API_KEY\"] = \"\"  # Set your Gemini API key here.\n",
    "# os.environ[\"ANTHROPIC_API_KEY\"] = \"\"  # Set your Anthropic API key here.\n",
    "\n",
    "# Set your GCP project id and region, if you want to query Llama 3.1 405B in Vertex.\n",
    "REGION = \"\"  # Set your GCP region here.\n",
    "PROJECT_ID = \"\"  # Set your GCP project id here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tutorial Directory\n",
    "\n",
    "Finally, we'll set up a directory to use for this tutorial, and some environment variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "tutorial_dir = \"letter_counting_tutorial\"\n",
    "\n",
    "Path(tutorial_dir).mkdir(parents=True, exist_ok=True)\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"  # Disable warnings from HF.\n",
    "\n",
    "# This is needed for vLLM to use multiple GPUs in a notebook.\n",
    "# If you're not running in a notebook, you can ignore this.\n",
    "os.environ[\"VLLM_WORKER_MULTIPROC_METHOD\"] = \"spawn\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "The dataset we'll use for this notebook is `oumi-ai/oumi-letter-count`, which can be found on [HF Datasets](https://huggingface.co/datasets/oumi-ai/oumi-letter-count). Its prompts ask to count the letters in various English words, with metadata in each example containing the correct count. We use the `train` split for training and the `test` split for evaluation. We'll use an Oumi dataset class, `LetterCountGrpoDataset`, to load and preprocess the HF Dataset. The following code displays an example prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2025-05-29 12:45:15,104][oumi][rank0][pid:5616][MainThread][INFO]][base_map_dataset.py:91] Creating map dataset (type: LetterCountGrpoDataset)... dataset_name: 'oumi-ai/oumi-letter-count'\n",
      "[2025-05-29 12:45:18,750][oumi][rank0][pid:5616][MainThread][INFO]][base_map_dataset.py:487] Dataset Info:\n",
      "\tSplit: validation\n",
      "\tVersion: 0.0.0\n",
      "\tDataset size: 22894322\n",
      "\tDownload size: 5697295\n",
      "\tSize: 28591617 bytes\n",
      "\tRows: 10000\n",
      "\tColumns: ['conversation_id', 'messages', 'metadata']\n",
      "[2025-05-29 12:45:19,064][oumi][rank0][pid:5616][MainThread][INFO]][base_map_dataset.py:426] Loaded DataFrame with shape: (10000, 3). Columns:\n",
      "conversation_id    object\n",
      "messages           object\n",
      "metadata           object\n",
      "dtype: object\n",
      "--------------------------------------------------------------------------------\n",
      "Sample:\n",
      "{'conversation_id': 'oumi_letter_count_0',\n",
      " 'messages': [{'content': 'Your final answer should be an integer written as '\n",
      "                          'digits and formatted as \"\\\\boxed{your_answer}\". For '\n",
      "                          'example, if the answer is 42, you should output '\n",
      "                          '\"\\\\boxed{42}\".',\n",
      "               'role': 'system'},\n",
      "              {'content': \"Could you determine the count of 'l's in \"\n",
      "                          \"'substantial'?\",\n",
      "               'role': 'user'}],\n",
      " 'metadata': {'letter': 'l',\n",
      "              'letter_count_integer': 1,\n",
      "              'letter_count_string': 'one',\n",
      "              'unformatted_prompt': 'Could you determine the count of '\n",
      "                                    '{letter}s in {word}?',\n",
      "              'word': 'substantial'}}\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "from oumi.datasets.grpo.letter_count import LetterCountGrpoDataset\n",
    "\n",
    "dataset = LetterCountGrpoDataset(\n",
    "    dataset=\"oumi-ai/oumi-letter-count-clean\", split=\"validation\"\n",
    ")\n",
    "print(\"-\" * 80)\n",
    "print(\"Sample:\")\n",
    "pprint(dataset.conversation(0).to_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation\n",
    "\n",
    "First, we'll evaluate how various models perform on the letter counting task. We'll evaluate frontier models by calling their respective remote API, and Llama 3.2 3B by running local inference on it using vLLM.\n",
    "\n",
    "We've already defined a custom evaluation function in Oumi which runs inference on the above dataset, extracts the answer from the model response, and calculates various metrics such as accuracy. This function is defined at `src/oumi/evaluation/registry/count_letters_task.py` ([GitHub link](https://github.com/oumi-ai/oumi/blob/main/src/oumi/evaluation/registry/count_letters_task.py)), and we print its contents below for reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def count_letters(\n",
      "    task_params: EvaluationTaskParams,\n",
      "    inference_engine: BaseInferenceEngine,\n",
      ") -> dict[str, Any]:\n",
      "    \"\"\"Custom evaluation function registered as `count_letters`.\"\"\"\n",
      "    dataset = LetterCountGrpoDataset(dataset=\"oumi-ai/oumi-letter-count-clean\", split=\"test\")\n",
      "    # TODO: OPE-1155: Add support for using Oumi dataset code to create the dataset.\n",
      "    # dataset = build_dataset(\"oumi-ai/oumi-letter-count\", tokenizer=None, sample_count=10)  # noqa: E501\n",
      "    num_samples = task_params.num_samples\n",
      "    if num_samples is None:\n",
      "        num_samples = len(dataset)\n",
      "    input_conversations = [dataset.conversation(i) for i in range(num_samples)]\n",
      "    conversations = inference_engine.infer(input_conversations)\n",
      "    logger.info(f\"Finished inference on {len(conversations)} conversations!\")\n",
      "    if len(conversations) > 0:\n",
      "        logger.info(f\"Sample conversation: {conversations[0]}\")\n",
      "\n",
      "    count = 0  # The number of examples with correct answers extracted.\n",
      "    total = 0  # All examples.\n",
      "    valid_count = 0  # The number of examples with valid answers extracted.\n",
      "    for i, conversation in enumerate(conversations):\n",
      "        total += 1\n",
      "        # Grab the model's response\n",
      "        response = conversation.last_message()\n",
      "        # Ignore cases where model didn't respond or it's a multimodal response.\n",
      "        # For now, we focus on text-only responses.\n",
      "        if not response or not isinstance(response.content, str):\n",
      "            continue\n",
      "        # Count the example as correct if the extracted prediction is correct.\n",
      "        prediction = _extract_prediction(response.content)\n",
      "        if prediction is None:\n",
      "            continue\n",
      "        valid_count += 1\n",
      "        if prediction == conversation.metadata[\"letter_count_integer\"]:\n",
      "            count += 1\n",
      "\n",
      "    return {\n",
      "        # Accuracy across all examples.\n",
      "        \"accuracy\": count / total if total > 0 else 0,\n",
      "        # Accuracy when only counting examples with properly extracted answers.\n",
      "        \"properly_extracted_accuracy\": count / valid_count if valid_count > 0 else 0,\n",
      "        \"num_samples\": num_samples,\n",
      "        # These three values sum up to num_samples.\n",
      "        \"num_correct_answers\": count,\n",
      "        \"num_incorrect_answers\": valid_count - count,\n",
      "        \"num_invalid_answers\": total - valid_count,\n",
      "    }\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import inspect\n",
    "\n",
    "from oumi.evaluation.registry.count_letters_task import count_letters\n",
    "\n",
    "print(inspect.getsource(count_letters))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following section, you can select which models you want to evaluate. You can lower `NUM_SAMPLES`  to reduce cost when calling remote APIs, with the downside of noisier results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_SAMPLES = 100\n",
    "# We set an environment variable to be used at the end of the Colab.\n",
    "os.environ[\"NUM_SAMPLES\"] = str(NUM_SAMPLES)\n",
    "\n",
    "model_names = [\n",
    "    \"llama_3b\",\n",
    "    # Uncomment any models you wish to evaluate - you can evaluate multiple at once.\n",
    "    # \"gpt_4o\",\n",
    "    # \"gemini_pro\",\n",
    "    # \"llama_405b\",\n",
    "    # \"claude_sonnet\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "❗**NOTICE:** If running this notebook on Colab, delete the following line: `inference_engine: VLLM`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting letter_counting_tutorial/llama_3b_eval.yaml\n"
     ]
    }
   ],
   "source": [
    "%%writefile $tutorial_dir/llama_3b_eval.yaml\n",
    "\n",
    "# We save this config as a YAML file as we'll use it again at the end of the notebook.\n",
    "model:\n",
    "  model_name: \"meta-llama/Llama-3.2-3B-Instruct\"\n",
    "  model_max_length: 131072\n",
    "  torch_dtype_str: \"bfloat16\"\n",
    "  attn_implementation: \"sdpa\"\n",
    "  trust_remote_code: True\n",
    "\n",
    "inference_engine: VLLM\n",
    "\n",
    "generation:\n",
    "  max_new_tokens: 2048\n",
    "\n",
    "tasks:\n",
    "  - evaluation_backend: custom\n",
    "    task_name: count_letters\n",
    "\n",
    "output_dir: \"letter_counting_tutorial/evaluation/llama_3b\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# EvaluationConfig for various models.\n",
    "# Note that Llama 3B uses the local VLLM inference engines, while the others use various\n",
    "# remote engines.\n",
    "\n",
    "with open(f\"{tutorial_dir}/llama_3b_eval.yaml\") as f:\n",
    "    llama_3b_yaml = f.read()\n",
    "\n",
    "configs = {\n",
    "    \"llama_3b\": llama_3b_yaml,\n",
    "    \"gpt_4o\": \"\"\"\n",
    "      model:\n",
    "        model_name: \"gpt-4o\"\n",
    "\n",
    "      inference_engine: OPENAI\n",
    "\n",
    "      inference_remote_params:\n",
    "        api_key_env_varname: \"OPENAI_API_KEY\"\n",
    "        max_retries: 3\n",
    "        num_workers: 100\n",
    "        politeness_policy: 60\n",
    "        connection_timeout: 300\n",
    "\n",
    "      generation:\n",
    "        max_new_tokens: 8192\n",
    "        temperature: 0.0\n",
    "\n",
    "      tasks:\n",
    "        - evaluation_backend: custom\n",
    "          task_name: count_letters\n",
    "\n",
    "      output_dir: \"letter_counting_tutorial/evaluation/gpt_4o\"\n",
    "      \"\"\",\n",
    "    \"gemini_pro\": \"\"\"\n",
    "      model:\n",
    "        model_name: \"gemini-2.5-pro-preview-03-25\"\n",
    "\n",
    "      inference_engine: GOOGLE_GEMINI\n",
    "\n",
    "      inference_remote_params:\n",
    "        api_key_env_varname: \"GEMINI_API_KEY\"\n",
    "        max_retries: 3\n",
    "        num_workers: 2\n",
    "        politeness_policy: 60\n",
    "        connection_timeout: 300\n",
    "\n",
    "      generation:\n",
    "        max_new_tokens: 8192\n",
    "        temperature: 0.0\n",
    "\n",
    "      tasks:\n",
    "        - evaluation_backend: custom\n",
    "          task_name: count_letters\n",
    "\n",
    "      output_dir: \"letter_counting_tutorial/evaluation/gemini_pro\"\n",
    "      \"\"\",\n",
    "    \"llama_405b\": f\"\"\"\n",
    "      model:\n",
    "        model_name: \"meta/llama-3.1-405b-instruct-maas\"\n",
    "\n",
    "      inference_engine: GOOGLE_VERTEX\n",
    "\n",
    "      inference_remote_params:\n",
    "        api_url: \"https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/openapi/chat/completions\"\n",
    "        max_retries: 3\n",
    "        num_workers: 10\n",
    "        politeness_policy: 60\n",
    "        connection_timeout: 300\n",
    "\n",
    "      generation:\n",
    "        max_new_tokens: 8192\n",
    "        temperature: 0.0\n",
    "\n",
    "      tasks:\n",
    "        - evaluation_backend: custom\n",
    "          task_name: count_letters\n",
    "\n",
    "      output_dir: \"letter_counting_tutorial/evaluation/llama_405b\"\n",
    "      \"\"\",\n",
    "    \"claude_sonnet\": \"\"\"\n",
    "      model:\n",
    "        model_name: \"claude-3-7-sonnet-latest\"\n",
    "\n",
    "      inference_engine: ANTHROPIC\n",
    "\n",
    "      inference_remote_params:\n",
    "        api_key_env_varname: \"ANTHROPIC_API_KEY\"\n",
    "        max_retries: 3\n",
    "        num_workers: 5\n",
    "        politeness_policy: 65\n",
    "        connection_timeout: 300\n",
    "\n",
    "      generation:\n",
    "        max_new_tokens: 8192\n",
    "        temperature: 0.0\n",
    "\n",
    "      tasks:\n",
    "        - evaluation_backend: custom\n",
    "          task_name: count_letters\n",
    "\n",
    "      output_dir: \"letter_counting_tutorial/evaluation/claude_sonnet\"\n",
    "      \"\"\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2025-05-23 14:15:06,778][oumi][rank0][pid:5616][MainThread][WARNING]][models.py:439] Undefined pad token. Setting it to `<|finetune_right_pad_id|>`.\n",
      "[2025-05-23 14:15:06,780][oumi][rank0][pid:5616][MainThread][INFO]][models.py:482] Using the model's built-in chat template for model 'meta-llama/Llama-3.2-3B-Instruct'.\n",
      "INFO 05-23 14:15:06 __init__.py:207] Automatically detected platform cuda.\n",
      "INFO 05-23 14:15:15 config.py:549] This model supports multiple tasks: {'score', 'reward', 'generate', 'embed', 'classify'}. Defaulting to 'generate'.\n",
      "WARNING 05-23 14:15:15 arg_utils.py:1187] Chunked prefill is enabled by default for models with max_model_len > 32K. Currently, chunked prefill might not work with some features or models. If you encounter any issues, please disable chunked prefill by setting --enable-chunked-prefill=False.\n",
      "INFO 05-23 14:15:15 config.py:1555] Chunked prefill is enabled with max_num_batched_tokens=2048.\n",
      "WARNING 05-23 14:15:15 cuda.py:95] To see benefits of async output processing, enable CUDA graph. Since, enforce-eager is enabled, async output processor cannot be used\n",
      "WARNING 05-23 14:15:15 config.py:685] Async output processing is not supported on the current platform type cuda.\n",
      "INFO 05-23 14:15:15 llm_engine.py:234] Initializing a V0 LLM engine (v0.7.3) with config: model='meta-llama/Llama-3.2-3B-Instruct', speculative_config=None, tokenizer='meta-llama/Llama-3.2-3B-Instruct', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.bfloat16, max_seq_len=131072, download_dir=None, load_format=LoadFormat.AUTO, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=True, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar'), observability_config=ObservabilityConfig(otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=meta-llama/Llama-3.2-3B-Instruct, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=True, chunked_prefill_enabled=True, use_async_output_proc=False, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={\"splitting_ops\":[],\"compile_sizes\":[],\"cudagraph_capture_sizes\":[],\"max_capture_size\":0}, use_cached_outputs=False, \n",
      "INFO 05-23 14:15:17 cuda.py:229] Using Flash Attention backend.\n",
      "INFO 05-23 14:15:18 model_runner.py:1110] Starting to load model meta-llama/Llama-3.2-3B-Instruct...\n",
      "INFO 05-23 14:15:19 weight_utils.py:254] Using model weights format ['*.safetensors']\n",
      "INFO 05-23 14:15:19 weight_utils.py:270] Time spent downloading weights for meta-llama/Llama-3.2-3B-Instruct: 0.556193 seconds\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "624ba55f2f4349dbaeb906b604c9e13a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading safetensors checkpoint shards:   0% Completed | 0/2 [00:00<?, ?it/s]\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 05-23 14:15:21 model_runner.py:1115] Loading model weights took 6.0160 GB\n",
      "INFO 05-23 14:15:22 worker.py:267] Memory profiling takes 0.66 seconds\n",
      "INFO 05-23 14:15:22 worker.py:267] the current vLLM instance can use total_gpu_memory (39.49GiB) x gpu_memory_utilization (0.90) = 35.55GiB\n",
      "INFO 05-23 14:15:22 worker.py:267] model weights take 6.02GiB; non_torch_memory takes 0.09GiB; PyTorch activation peak memory takes 1.18GiB; the rest of the memory reserved for KV Cache is 28.25GiB.\n",
      "INFO 05-23 14:15:22 executor_base.py:111] # cuda blocks: 16532, # CPU blocks: 2340\n",
      "INFO 05-23 14:15:22 executor_base.py:116] Maximum concurrency for 131072 tokens per request: 2.02x\n",
      "INFO 05-23 14:15:28 llm_engine.py:436] init engine (profile, create kv cache, warmup model) took 7.09 seconds\n",
      "[2025-05-23 14:15:29,029][oumi][rank0][pid:5616][MainThread][INFO]][base_map_dataset.py:91] Creating map dataset (type: LetterCountCleanGrpoDataset)... dataset_name: 'oumi-ai/oumi-letter-count-clean'\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d2b4916eb4f84956a2a94e68e57b6c16",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "train-00000-of-00001.parquet:   0%|          | 0.00/4.47M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "27b5e9ad11e346f6a614268a19938d14",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "validation-00000-of-00001.parquet:   0%|          | 0.00/400k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc24bd85d897442daf730b27265f21f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "test-00000-of-00001.parquet:   0%|          | 0.00/831k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f6d90ea6b03343f1b8013fdb01cb6867",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/100000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "18d87f39517045b3ab7f42b09d440d4e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating validation split:   0%|          | 0/10000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "69fc5aab14a84ba3aceac062a20f08b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating test split:   0%|          | 0/19999 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2025-05-23 14:15:36,180][oumi][rank0][pid:5616][MainThread][INFO]][base_map_dataset.py:487] Dataset Info:\n",
      "\tSplit: test\n",
      "\tVersion: 0.0.0\n",
      "\tDataset size: 22894154\n",
      "\tDownload size: 5697150\n",
      "\tSize: 28591304 bytes\n",
      "\tRows: 19999\n",
      "\tColumns: ['conversation_id', 'messages', 'metadata']\n",
      "[2025-05-23 14:15:36,563][oumi][rank0][pid:5616][MainThread][INFO]][base_map_dataset.py:426] Loaded DataFrame with shape: (19999, 3). Columns:\n",
      "conversation_id    object\n",
      "messages           object\n",
      "metadata           object\n",
      "dtype: object\n",
      "INFO 05-23 14:15:36 chat_utils.py:332] Detected the chat template content format to be 'string'. You can set `--chat-template-content-format` to override this.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processed prompts: 100%|██████████| 100/100 [00:02<00:00, 43.31it/s, est. speed input: 3769.97 toks/s, output: 2047.43 toks/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2025-05-23 14:15:38,950][oumi][rank0][pid:5616][MainThread][INFO]][count_letters_task.py:52] Finished inference on 100 conversations!\n",
      "[2025-05-23 14:15:38,951][oumi][rank0][pid:5616][MainThread][INFO]][count_letters_task.py:54] Sample conversation: conversation_id='oumi_letter_count_1' messages=[SYSTEM: Your final answer should be an integer written as digits and formatted as \"\\boxed{your_answer}\". For example, if the answer is 42, you should output \"\\boxed{42}\"., USER: Look through 'drumwood' and count the 'e's., ASSISTANT: There are 1 'e's in 'drumwood'.] metadata={'letter': 'e', 'letter_count_integer': 0, 'letter_count_string': 'zero', 'unformatted_prompt': 'Look through {word} and count the {letter}s.', 'word': 'drumwood'}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Run evaluation on all specified models.\n",
    "\n",
    "from oumi.core.configs import EvaluationConfig\n",
    "from oumi.core.evaluation import Evaluator\n",
    "\n",
    "results = {}\n",
    "\n",
    "for model_name in model_names:\n",
    "    # Create the evaluation config from the YAML string.\n",
    "    config_yaml: str = configs[model_name]\n",
    "    config = EvaluationConfig.from_str(config_yaml)\n",
    "    config.tasks[0].num_samples = NUM_SAMPLES\n",
    "\n",
    "    # Run the evaluation.\n",
    "    evaluator = Evaluator()\n",
    "    evaluator_out = evaluator.evaluate(config)\n",
    "\n",
    "    # # Record the results.\n",
    "    results[model_name] = evaluator_out[0].get_results()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total samples: 100\n",
      "--------------------------------------------------------------------------------\n",
      "Model: llama_3b\n",
      "Accuracy: 31.00%\n",
      "Properly Extracted Accuracy: 46.27%\n",
      "Num correct, incorrect, invalid: 31, 36, 33\n"
     ]
    }
   ],
   "source": [
    "# Print results.\n",
    "\n",
    "print(f\"Total samples: {NUM_SAMPLES}\")\n",
    "for model_name, result in results.items():\n",
    "    print(\"-\" * 80)\n",
    "    print(f\"Model: {model_name}\")\n",
    "    print(f\"Accuracy: {result['accuracy']:.2%}\")\n",
    "    print(f\"Properly Extracted Accuracy: {result['properly_extracted_accuracy']:.2%}\")\n",
    "    correct = result[\"num_correct_answers\"]\n",
    "    incorrect = result[\"num_incorrect_answers\"]\n",
    "    invalid = result[\"num_invalid_answers\"]\n",
    "    print(f\"Num correct, incorrect, invalid: {correct}, {incorrect}, {invalid}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GRPO\n",
    "\n",
    "Now, we train Llama 3.2 3B on the task of counting letters using the GRPO algorithm implemented by [HuggingFace's `trl` library](https://huggingface.co/docs/trl/en/index).\n",
    "\n",
    "Note that we can calculate a concrete reward for this task by comparing the answer extracted by the model with the correct answer. In the reward function defined in `src/oumi/datasets/grpo/rewards/count_letters_rewards.py` ([GitHub link](https://github.com/oumi-ai/oumi/blob/main/src/oumi/datasets/grpo/rewards/count_letters_rewards.py)), we calculate the reward to be `-abs(predicted_count - target_count)`. We use simple heuristics to extract the predicted count. The following cell prints out the reward function code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Copyright 2025 - Oumi\n",
      "#\n",
      "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
      "# you may not use this file except in compliance with the License.\n",
      "# You may obtain a copy of the License at\n",
      "#\n",
      "#     http://www.apache.org/licenses/LICENSE-2.0\n",
      "#\n",
      "# Unless required by applicable law or agreed to in writing, software\n",
      "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
      "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
      "# See the License for the specific language governing permissions and\n",
      "# limitations under the License.\n",
      "\n",
      "import re\n",
      "from typing import Any, Optional\n",
      "\n",
      "from oumi.core.registry import RegistryType, register\n",
      "\n",
      "\n",
      "def _extract_prediction(response: str) -> Optional[int]:\n",
      "    r\"\"\"Returns the numeric answer extracted from `\\boxed{...}`, or None otherwise.\"\"\"\n",
      "    regex_result = re.findall(r\"\\\\boxed\\{([-+]?\\d+)\\}\", response)\n",
      "    if not regex_result or len(regex_result) != 1:\n",
      "        return None\n",
      "    number_str = regex_result[0]\n",
      "    # Except clause shouldn't trigger because the regex should only find ints.\n",
      "    try:\n",
      "        return int(number_str)\n",
      "    except ValueError:\n",
      "        return None\n",
      "\n",
      "\n",
      "def compute_letter_count_reward(completion: str, target_count: int) -> float:\n",
      "    \"\"\"Computes the rewards for counting the letters in a string.\n",
      "\n",
      "    Args:\n",
      "        completion: The completion string from the LLM.\n",
      "        target_count: The target count of letters.\n",
      "\n",
      "    Returns:\n",
      "        The reward value.\n",
      "    \"\"\"\n",
      "    count = _extract_prediction(completion)\n",
      "\n",
      "    # Lowest reward goes to unparseable responses\n",
      "    if count is None:\n",
      "        return -3.0\n",
      "\n",
      "    delta = abs(count - target_count)\n",
      "\n",
      "    # Reward scales from [0, -2) as delta increases\n",
      "    # Ensures that \"worse\" answers (where the counts are off by a higher amount) are\n",
      "    # penalized while never reaching -3.0 which is reserved for unparseable answers.\n",
      "    return (1 / (delta + 0.5)) - 2\n",
      "\n",
      "\n",
      "@register(\"count_letters\", RegistryType.REWARD_FUNCTION)\n",
      "def _count_letters(\n",
      "    completions: list[list[dict[str, Any]]],\n",
      "    letter_count: list[int],\n",
      "    **kwargs: dict[str, Any],\n",
      ") -> list[float]:\n",
      "    \"\"\"Custom reward function for counting letters in a string.\n",
      "\n",
      "    For more details on custom reward functions used in trl's GRPOTrainer, see:\n",
      "    https://huggingface.co/docs/trl/main/en/grpo_trainer#using-a-custom-reward-function.\n",
      "\n",
      "    Args:\n",
      "        completions: The list of completions from the LLM.\n",
      "        letter_count: The list of target count of letters.\n",
      "        kwargs: Unused.\n",
      "\n",
      "    Returns:\n",
      "        The reward values for each completion, calculated as the negative of the\n",
      "        absolute difference between the count and the target count. The count is assumed\n",
      "        to be the last group of consecutive digits in the completion string.\n",
      "    \"\"\"\n",
      "    completions_strs = [c[0][\"content\"] for c in completions]\n",
      "    return [\n",
      "        compute_letter_count_reward(c, t)\n",
      "        for c, t in zip(completions_strs, letter_count)\n",
      "    ]\n"
     ]
    }
   ],
   "source": [
    "!cat ../src/oumi/datasets/grpo/rewards/count_letters_rewards.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up to free-up GPU memory used for evaluation above\n",
    "import gc\n",
    "\n",
    "import torch\n",
    "\n",
    "\n",
    "def cleanup_memory():\n",
    "    \"\"\"Delete the evaluator and collect garbage.\"\"\"\n",
    "    global evaluator\n",
    "    if evaluator:  # type: ignore\n",
    "        del evaluator\n",
    "        evaluator = None\n",
    "    for _ in range(3):\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "        torch.cuda.synchronize()\n",
    "\n",
    "\n",
    "cleanup_memory()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "❗**NOTICE:** Set `training.enable_wandb` to True if you want to log your training run to Weights and Biases. In addition, you must also log into WandB, ex. by running `wandb login`.\n",
    "\n",
    "❗**NOTICE:** We only train for 2 steps for demonstration purposes. You can increase `max_steps`, or replace it with `num_train_epochs` to set your desired number of epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing letter_counting_tutorial/grpo_train.yaml\n"
     ]
    }
   ],
   "source": [
    "%%writefile $tutorial_dir/grpo_train.yaml\n",
    "\n",
    "model:\n",
    "  model_name: \"meta-llama/Llama-3.2-3B-Instruct\"\n",
    "  model_max_length: 8192\n",
    "  torch_dtype_str: \"bfloat16\"\n",
    "  attn_implementation: \"sdpa\"\n",
    "\n",
    "data:\n",
    "  train:\n",
    "    datasets:\n",
    "      - dataset_name: \"oumi-ai/oumi-letter-count\"\n",
    "        split: \"train\"\n",
    "\n",
    "training:\n",
    "  trainer_type: \"TRL_GRPO\"\n",
    "  save_steps: 500\n",
    "  # max_steps: 500\n",
    "  per_device_train_batch_size: 2\n",
    "  gradient_accumulation_steps: 1\n",
    "  learning_rate: 5e-7\n",
    "  lr_scheduler_type: \"cosine\"\n",
    "  warmup_steps: 20\n",
    "\n",
    "  reward_functions: [\"count_letters\"]\n",
    "\n",
    "  ddp_find_unused_parameters: False\n",
    "  optimizer: \"adafactor\"\n",
    "  compile: True\n",
    "\n",
    "  grpo:\n",
    "    num_generations: 4\n",
    "\n",
    "  dataloader_num_workers: \"auto\"\n",
    "  dataloader_prefetch_factor: 32\n",
    "\n",
    "  logging_steps: 10\n",
    "  output_dir: \"letter_counting_tutorial/llama_3b_grpo\"\n",
    "  # Set this to True if you want to log to Weights and Biases.\n",
    "  enable_wandb: False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2025-04-10 09:38:32,770][oumi][rank0][pid:10890][MainThread][INFO]][distributed_run.py:276] Running the command: ['torchrun', '--nnodes=1', '--node-rank=0', '--nproc-per-node=4', '--master-addr=127.0.0.1', '--master-port=8007', '-m', 'oumi', 'train', '-c', 'letter_counting_tutorial/grpo_train.yaml']\n",
      "\n",
      "\u001b[32m   ____  _    _ __  __ _____\u001b[0m\n",
      "\u001b[32m  / __ \\| |  | |  \\/  |_   _|\u001b[0m\n",
      "\u001b[32m | |  | | |  | | \\  / | | |\u001b[0m\n",
      "\u001b[32m | |  | | |  | | |\\/| | | |\u001b[0m\n",
      "\u001b[32m | |__| | |__| | |  | |_| |_\u001b[0m\n",
      "\u001b[32m  \\____/ \\____/|_|  |_|_____|\u001b[0m\n",
      "\u001b[2K\u001b[32m⠦\u001b[0m \u001b[32mLoading configuration...\u001b[0m\n",
      "Model Parameters Summary:\n",
      "🔢 Total     parameters: 3,212,749,824\n",
      "🔗 Embedding parameters: 394,002,432\n",
      "🎯 Trainable parameters: 3,212,749,824\n",
      "🔒 Frozen    parameters: 0 (0.00%)\n",
      "\n",
      "[2025-04-10 09:38:47,784][oumi][rank0][pid:10894][MainThread][INFO]][base_map_dataset.py:91] Creating map dataset (type: LetterCountGrpoDataset)... dataset_name: 'oumi-ai/oumi-letter-count'\n",
      "[2025-04-10 09:38:50,326][oumi][rank0][pid:10894][MainThread][INFO]][base_map_dataset.py:487] Dataset Info:\n",
      "\tSplit: train\n",
      "\tVersion: 0.0.0\n",
      "\tDataset size: 22894322\n",
      "\tDownload size: 5697295\n",
      "\tSize: 28591617 bytes\n",
      "\tRows: 100000\n",
      "\tColumns: ['conversation_id', 'messages', 'metadata']\n",
      "Generating train split: 8 examples [00:00, 2656.51 examples/s]\n",
      "[2025-04-10 09:38:50,827][oumi][rank0][pid:10894][MainThread][INFO]][base_map_dataset.py:426] Loaded DataFrame with shape: (100000, 3). Columns:\n",
      "conversation_id    object\n",
      "messages           object\n",
      "metadata           object\n",
      "dtype: object\n",
      "[2025-04-10 09:38:50,837][oumi][rank0][pid:10894][MainThread][INFO]][base_map_dataset.py:312] LetterCountGrpoDataset: features=dict_keys(['prompt', 'letter_count'])\n",
      "Generating train split: 100000 examples [00:08, 11203.44 examples/s]\n",
      "[2025-04-10 09:39:12,603][oumi][rank0][pid:10894][MainThread][INFO]][base_map_dataset.py:376] Finished transforming dataset (LetterCountGrpoDataset)! Speed: 4594.28 examples/sec. Examples: 100000. Duration: 21.8 sec. Transform workers: 1.\n",
      "[2025-04-10 09:39:14,799][oumi][rank0][pid:10894][MainThread][INFO]][train.py:419] Training init time: 31.849s\n",
      "[2025-04-10 09:39:14,800][oumi][rank0][pid:10894][MainThread][INFO]][train.py:420] Starting training... (TrainerType.TRL_GRPO, transformers: 4.51.1)\n",
      "{'train_runtime': 226.3313, 'train_samples_per_second': 0.071, 'train_steps_per_second': 0.009, 'train_loss': -0.5479733943939209, 'rewards/_count_letters': -1.0625, 'reward': -1.0625, 'reward_std': 0.9289332032203674, 'completion_length': 47.8125, 'kl': 0.0001125335693359375, 'epoch': 0.0}\n",
      "100%|████████████████████████████████████████████| 2/2 [03:46<00:00, 113.16s/it]\n",
      "[2025-04-10 09:43:01,562][oumi][rank0][pid:10894][MainThread][INFO]][train.py:427] Training is Complete.\n",
      "[2025-04-10 09:43:01,562][oumi][rank0][pid:10894][MainThread][INFO]][device_utils.py:297] GPU Metrics After Training: GPU runtime info: None.\n",
      "[2025-04-10 09:43:01,562][oumi][rank0][pid:10894][MainThread][INFO]][torch_utils.py:136] Peak GPU memory usage: 29.84 GB\n",
      "[2025-04-10 09:43:01,562][oumi][rank0][pid:10894][MainThread][INFO]][train.py:434] Saving final state...\n",
      "[2025-04-10 09:43:01,563][oumi][rank0][pid:10894][MainThread][INFO]][train.py:439] Saving final model...\n",
      "[2025-04-10 09:43:14,924][oumi][rank0][pid:10894][MainThread][INFO]][hf_trainer.py:116] Model has been saved at letter_counting_tutorial/llama_3b_grpo\n",
      "[2025-04-10 09:43:15,458][oumi][rank0][pid:10894][MainThread][INFO]][train.py:446] \n",
      "\n",
      "» We're always looking for feedback. What's one thing we can improve? https://oumi.ai/feedback\n",
      "[2025-04-10 09:43:23,534][oumi][rank0][pid:10890][MainThread][INFO]][distributed_run.py:295] Successfully completed! (Rank: 0. Duration: 290.8 sec)\n"
     ]
    }
   ],
   "source": [
    "!oumi distributed torchrun -m oumi train -c $tutorial_dir/grpo_train.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating our Trained Model\n",
    "\n",
    "Let's now evaluate our trained model to see if it improved on the letter counting task. Note that it may not improve much, since we trained it for a relatively short time.\n",
    "\n",
    "Below, we demonstrate an alternative method of running evaluation with the `oumi` CLI. We use the same Llama 3B evaluation config we used above, with the only change being pointing it at the model we just trained.\n",
    "\n",
    "First, we need to reset the notebook to clear variables from our previous vLLM run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reset -f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[32m   ____  _    _ __  __ _____\u001b[0m\n",
      "\u001b[32m  / __ \\| |  | |  \\/  |_   _|\u001b[0m\n",
      "\u001b[32m | |  | | |  | | \\  / | | |\u001b[0m\n",
      "\u001b[32m | |  | | |  | | |\\/| | | |\u001b[0m\n",
      "\u001b[32m | |__| | |__| | |  | |_| |_\u001b[0m\n",
      "\u001b[32m  \\____/ \\____/|_|  |_|_____|\u001b[0m\n",
      "\u001b[2K\u001b[32m⠴\u001b[0m \u001b[32mLoading configuration...\u001b[0m0m\n",
      "\u001b[2K\u001b[32m⠋\u001b[0m \u001b[32mRunning evaluation...\u001b[0m[2025-04-10 09:47:15,521][oumi][rank0][pid:16694][MainThread][INFO]][models.py:482] Using the model's built-in chat template for model 'letter_counting_tutorial/llama_3b_grpo'.\n",
      "\u001b[2KINFO 04-10 09:47:15 __init__.py:207] Automatically detected platform cuda.\n",
      "\u001b[2KINFO 04-10 09:47:33 model_runner.py:1110] Starting to load model \n",
      "letter_counting_tutorial/llama_3b_grpo...\n",
      "\u001b[32m⠸\u001b[0m \u001b[32mRunning evaluation...\u001b[0m\u001b[1;36m(VllmWorkerProcess pid=16718)\u001b[0;0m INFO 04-10 09:47:33 model_runner.py:1110] Starting to load model letter_counting_tutorial/llama_3b_grpo...\n",
      "\u001b[32m⠴\u001b[0m \u001b[32mRunning evaluation...\u001b[0m\u001b[1;36m(VllmWorkerProcess pid=16719)\u001b[0;0m INFO 04-10 09:47:34 model_runner.py:1115] Loading model weights took 1.5341 GB\n",
      "\u001b[1;36m(VllmWorkerProcess pid=16718)\u001b[0;0m INFO 04-10 09:47:34 model_runner.py:1115] Loading model weights took 1.5341 GB\n",
      "\n",
      "\u001b[2K\u001b[32m⠹\u001b[0m \u001b[32mRunning evaluation...\u001b[0m[2025-04-10 09:47:46,962][oumi][rank0][pid:16694][MainThread][INFO]][base_map_dataset.py:91] Creating map dataset (type: LetterCountGrpoDataset)... dataset_name: 'oumi-ai/oumi-letter-count'\n",
      "\u001b[2K\u001b[32m⠋\u001b[0m \u001b[32mRunning evaluation...\u001b[0m[2025-04-10 09:47:49,180][oumi][rank0][pid:16694][MainThread][INFO]][base_map_dataset.py:487] Dataset Info:\n",
      "\tSplit: test\n",
      "\tVersion: 0.0.0\n",
      "\tDataset size: 22894322\n",
      "\tDownload size: 5697295\n",
      "\tSize: 28591617 bytes\n",
      "\tRows: 20000\n",
      "\tColumns: ['conversation_id', 'messages', 'metadata']\n",
      "\u001b[2K\u001b[32m⠏\u001b[0m \u001b[32mRunning evaluation...\u001b[0m[2025-04-10 09:47:49,820][oumi][rank0][pid:16694][MainThread][INFO]][base_map_dataset.py:426] Loaded DataFrame with shape: (20000, 3). Columns:\n",
      "conversation_id    object\n",
      "messages           object\n",
      "metadata           object\n",
      "dtype: object\n",
      "\u001b[2KINFO 04-10 09:47:49 chat_utils.py:332] Detected the chat template content format\n",
      "to be 'string'. You can set `--chat-template-content-format` to override this.\n",
      "\u001b[2KProcessed prompts: \u001b[1;36m100\u001b[0m%|#| \u001b[1;36m100\u001b[0m/\u001b[1;36m100\u001b[0m \u001b[1m[\u001b[0m\u001b[1;92m00:03\u001b[0m<\u001b[1;92m00:00\u001b[0m, \u001b[1;36m28.\u001b[0m84it/s, est. speed input: \u001b[1;36m26\u001b[0m\n",
      "\u001b[32m⠸\u001b[0m \u001b[32mRunning evaluation...\u001b[0m[2025-04-10 09:47:53,387][oumi][rank0][pid:16694][MainThread][INFO]][count_letters_task.py:53] Finished inference on 100 conversations!\n",
      "[2025-04-10 09:47:53,387][oumi][rank0][pid:16694][MainThread][INFO]][count_letters_task.py:55] Sample conversation: conversation_id='oumi_letter_count_0' messages=[USER: Look through 'perivaginal' and count the 'n's., SYSTEM: Your final answer should be written as digits and formatted as \"\\boxed{your_answer}\". For example, if the answer is 42, make sure to output \"\\boxed{42}\"., ASSISTANT: There are 2 'n's in 'perivaginal'. \n",
      "\n",
      "\\boxed{2}] metadata={'letter': 'n', 'letter_count_integer': 1, 'letter_count_string': 'one', 'unformatted_prompt': 'Look through {word} and count the {letter}s.', 'word': 'perivaginal'}\n",
      "\u001b[2KINFO 04-10 09:47:53 multiproc_worker_utils.py:141] Terminating local vLLM worker\n",
      "processes\n",
      "\u001b[32m⠧\u001b[0m \u001b[32mRunning evaluation...\u001b[0m\u001b[1;36m(VllmWorkerProcess pid=16717)\u001b[0;0m INFO 04-10 09:47:53 multiproc_worker_utils.py:253] Worker exiting\n",
      "\u001b[2K\u001b[32m⠧\u001b[0m \u001b[32mRunning evaluation...\u001b[0m\n",
      "\u001b[1A\u001b[2K\u001b[1;35m                         Evaluation Results                         \u001b[0m\n",
      "┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━┓\n",
      "┃\u001b[1m \u001b[0m\u001b[1mBenchmark    \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mMetric                     \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mScore \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mStd Error\u001b[0m\u001b[1m \u001b[0m┃\n",
      "┡━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━┩\n",
      "│\u001b[36m \u001b[0m\u001b[36mcount_letters\u001b[0m\u001b[36m \u001b[0m│\u001b[33m \u001b[0m\u001b[33mAccuracy                   \u001b[0m\u001b[33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m51.00%\u001b[0m\u001b[32m \u001b[0m│\u001b[2m \u001b[0m\u001b[2m-        \u001b[0m\u001b[2m \u001b[0m│\n",
      "├───────────────┼─────────────────────────────┼────────┼───────────┤\n",
      "│\u001b[36m \u001b[0m\u001b[36mcount_letters\u001b[0m\u001b[36m \u001b[0m│\u001b[33m \u001b[0m\u001b[33mProperly Extracted Accuracy\u001b[0m\u001b[33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m53.68%\u001b[0m\u001b[32m \u001b[0m│\u001b[2m \u001b[0m\u001b[2m-        \u001b[0m\u001b[2m \u001b[0m│\n",
      "├───────────────┼─────────────────────────────┼────────┼───────────┤\n",
      "│\u001b[36m \u001b[0m\u001b[36mcount_letters\u001b[0m\u001b[36m \u001b[0m│\u001b[33m \u001b[0m\u001b[33mNum Samples                \u001b[0m\u001b[33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m100   \u001b[0m\u001b[32m \u001b[0m│\u001b[2m \u001b[0m\u001b[2m-        \u001b[0m\u001b[2m \u001b[0m│\n",
      "├───────────────┼─────────────────────────────┼────────┼───────────┤\n",
      "│\u001b[36m \u001b[0m\u001b[36mcount_letters\u001b[0m\u001b[36m \u001b[0m│\u001b[33m \u001b[0m\u001b[33mNum Correct Answers        \u001b[0m\u001b[33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m51    \u001b[0m\u001b[32m \u001b[0m│\u001b[2m \u001b[0m\u001b[2m-        \u001b[0m\u001b[2m \u001b[0m│\n",
      "├───────────────┼─────────────────────────────┼────────┼───────────┤\n",
      "│\u001b[36m \u001b[0m\u001b[36mcount_letters\u001b[0m\u001b[36m \u001b[0m│\u001b[33m \u001b[0m\u001b[33mNum Incorrect Answers      \u001b[0m\u001b[33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m44    \u001b[0m\u001b[32m \u001b[0m│\u001b[2m \u001b[0m\u001b[2m-        \u001b[0m\u001b[2m \u001b[0m│\n",
      "├───────────────┼─────────────────────────────┼────────┼───────────┤\n",
      "│\u001b[36m \u001b[0m\u001b[36mcount_letters\u001b[0m\u001b[36m \u001b[0m│\u001b[33m \u001b[0m\u001b[33mNum Invalid Answers        \u001b[0m\u001b[33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m5     \u001b[0m\u001b[32m \u001b[0m│\u001b[2m \u001b[0m\u001b[2m-        \u001b[0m\u001b[2m \u001b[0m│\n",
      "└───────────────┴─────────────────────────────┴────────┴───────────┘\n",
      "[rank0]:[W410 09:47:58.432638348 ProcessGroupNCCL.cpp:1250] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present,  but this warning has only been added since PyTorch 2.4 (function operator())\n"
     ]
    }
   ],
   "source": [
    "!oumi evaluate -c letter_counting_tutorial/llama_3b_eval.yaml \\\n",
    "    --model.model_name \"letter_counting_tutorial/llama_3b_grpo\" \\\n",
    "    --tasks.0.num_samples $NUM_SAMPLES \\\n",
    "    --output_dir \"letter_counting_tutorial/evaluation/llama_3_grpo\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Better Letter Counter\n",
    "\n",
    "Looks like we were able to significantly improve on the performance of Llama-3.2-3B-Instruct:\n",
    "\n",
    "**BEFORE**\n",
    "\n",
    "Accuracy: 31.00%\n",
    "Properly Extracted Accuracy: 46.27%\n",
    "\n",
    "**AFTER**\n",
    "\n",
    "Accuracy: 51.00%\n",
    "Properly Extracted Accuracy: 53.68%\n",
    "\n",
    "A lot of the improvement from using GRPO came because this small LLM learned to better mimic the expected output format of the extractor, but the accuracy for properly extracted samples also improved! This is a great illustration of the kind of task GRPO training excels at."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's Next?\n",
    "\n",
    "Now that you know how easy it is to train in Oumi using GRPO, perhaps you'd like to try training on your own data (in a similar data format) -- check out [our docs](https://oumi.ai/docs/en/latest/resources/datasets/sft_datasets.html#using-an-unregistered-dataset-whose-format-is-identical-to-a-registered-dataset) for an easy way to do just that. \n",
    "\n",
    "Have fun!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.12 (myenv)",
   "language": "python",
   "name": "myenv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
